Require Import List.
Require Export FDPIA_build.
Require Import Strings.String.
Import ListNotations.
Open Scope string_scope.

Section product.

Local Open Scope string.

Variable n:nat. 
Hypothesis wtn:  writeNat n = "n".
Let xs : exp (ary '{n} num) := (vrd "xs").
Let ys : exp (ary '{n} num) := (vrd "ys").
Let zero : exp num := vrd "0".


Example product_par_code : exists output,
  CodeGen_comm ( A{reduce add zero 
   (mapPar (cur mul : exp(tensor num) -> exp num) (zip xs ys)), 
     acc' (Var "out") }) 0 = output.
Proof.
   Translate.
  exists "float v0[n];
#pragma omp parallel for
for (int i1=0; i1 < n; i1 += 1) 
  { v0[i1] = (xs[i1] * ys[i1]) ;}
float v2;
v2 = 0 ;
for(int i3=0; i3<n; i3 += 1) 
  { v2 = (v0[i3] + v2) ;}
out = v2 ;". reflexivity.
Qed.

End product.

Section product_parallel.

Local Open Scope string.

Variable n m p:nat. 
Hypothesis wtn:  writeNat n = "n".
Hypothesis wtm:  writeNat m = "m".
Hypothesis wtp:  writeNat p = "p".
Let xs : exp (ary '{n*(m*p)} num) := (vrd "xs").
Let ys : exp (ary '{n*(m*p)} num) := (vrd "ys").
Let zero : exp num := vrd "0".

Example product_code : exists output,
  CodeGen_comm ( A{reduce add zero (join
   (mapPar (fun zs1 => 
     mapPar (reduce (fun x a =>add (mul (Fst x) (Snd x)) a) zero)
      (split zs1)) (split (zip xs ys)))), 
     acc' (Var "out") }) 0 = output.
Proof.
  Translate.
  exists "float v0[n][m];
#pragma omp parallel for
for (int i1=0; i1 < n; i1 += 1) 
  { #pragma omp parallel for
for (int i2=0; i2 < m; i2 += 1) 
  { float v3;
v3 = 0 ;
for(int i4=0; i4<p; i4 += 1) 
  { v3 = ((xs[i1*m*p+i2*p+i4] * ys[i1*m*p+i2*p+i4]) + v3) ;}
v0[i1][i2] = v3 ;}}
float v5;
v5 = 0 ;
for(int i6=0; i6<n*m; i6 += 1) 
  { v5 = (v0[i6%n][i6/n] + v5) ;}
out = v5 ;". reflexivity.
Qed.

End product_parallel.


(**************************** vector ****************************)

(* product *)
Definition product_par {n} (zero:exp num)  (xs ys:exp(ary '{n} num)):= 
  reduce add zero (mapPar (cur mul) (zip xs ys)).

Definition product_seq zero {n} (xs ys:exp(ary '{n} num)):= 
   reduce add zero (mapSeq (cur mul) (zip xs ys)).

Definition product_split {m n:nat} (zero:exp num) (xs ys: exp (ary '{m*n} num)):=
  reduce add zero(mapPar (fun zs => 
    reduce (fun x a => add (add (Fst x) (Snd x)) a) zero zs)(split (zip xs ys))).

(* vector addition: a + b *)
Definition vec_add_seq {n} (xs ys: exp(ary '{n} num)):=
  mapSeq (cur add) (zip xs ys).

Definition vec_add_par {n} (xs ys: exp(ary '{n} num)):=
  mapPar (cur add) (zip xs ys).

(* vector × matrix : a * A *)
Definition l_mul_dl_ss zero {m n:nat} (dl:exp(ary '{m}(ary '{n} num))) (l:exp(ary '{n} num)) :=
  mapSeq (product_seq zero l) dl. 

Definition l_mul_dl_sp zero {m n:nat} (dl:exp(ary '{m}(ary '{n} num))) (l:exp(ary '{n} num)) :=
  mapPar (product_seq zero l) dl. 

Definition l_mul_dl_pp zero {m n:nat} (dl:exp(ary '{m}(ary '{n} num))) (l:exp(ary '{n} num)) :=
  mapPar (product_par zero l) dl. 

Definition l_mul_dl_ps zero {m n:nat} (dl:exp(ary '{m}(ary '{n} num))) (l:exp(ary '{n} num)) :=
  mapPar (product_seq zero l) dl.

(**************************** matrix ****************************)

(* matrix × matrix without transpose *)
Definition dl_mul_dl_sss zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapSeq (l_mul_dl_ss zero dl2) dl1.

Definition dl_mul_dl_ssp zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapSeq (l_mul_dl_sp zero dl2) dl1.

Definition dl_mul_dl_spp zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapSeq (l_mul_dl_pp zero dl2) dl1.

Definition dl_mul_dl_sps zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapSeq (l_mul_dl_ps zero dl2) dl1.

Definition dl_mul_dl_ppp zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapPar (l_mul_dl_pp zero dl2) dl1.

Definition dl_mul_dl_pps zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapPar (l_mul_dl_ps zero dl2) dl1.

Definition dl_mul_dl_psp zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapPar (l_mul_dl_sp zero dl2) dl1.

Definition dl_mul_dl_pss zero {m n p:nat} (dl1: exp(ary '{m} (ary '{n} num)))
  (dl2:exp(ary '{p} (ary '{n} num))) := mapPar (l_mul_dl_ss zero dl2) dl1.

(* matrix × vector *)
Definition dl_mul_l_ps zero {m n :nat} (dl: exp(ary '{m} (ary '{n} num)))
  (l: exp (ary '{n} num)) := mapPar (product_seq zero l) dl.

(* matrix addition *)
Definition mat_add_ps {m n :nat} (xs ys:exp(ary '{m}(ary '{n} num))) :=
  mapPar (cur (vec_add_seq)) (zip xs ys).

Definition mat_add_sp {m n :nat} (xs ys:exp(ary '{m}(ary '{n} num))) :=
  mapSeq (cur (vec_add_par)) (zip xs ys).

Definition mat_add_pp {m n :nat} (xs ys:exp(ary '{m}(ary '{n} num))) :=
  mapPar (cur (vec_add_par)) (zip xs ys).

Definition mat_add_ss {m n :nat} (xs ys:exp(ary '{m}(ary '{n} num))) :=
  mapSeq (cur (vec_add_seq)) (zip xs ys).

(**************************** block matrix ****************************)

Definition block_product {m n n1:nat} zero (mzero:exp (ary '{m} (ary '{m} num)))
    (xs :exp(ary '{n1} (ary '{m} (ary '{n} num))))
   (ys :exp(ary '{n1} (ary '{m} (ary '{n} num)))):=
  reduce mat_add_ps mzero
    (mapSeq (cur (dl_mul_dl_sss zero)) (zip xs ys)).

Definition block_l_mul_dl{m n m1 n1:nat} zero (mzero:exp (ary '{m} (ary '{m} num)))
    (dl:exp(ary '{m1} (ary '{n1} (ary '{m} (ary '{n} num)))))  (l:exp(ary '{n1} (ary '{m} (ary '{n} num)))):= 
  mapPar (block_product zero mzero l) dl. 

(* bA * bB *)
Definition block_dl_mul_dl{m n m1 n1 p1:nat}zero (mzero:exp (ary '{m} (ary '{m} num)))
    (dl1:exp(ary '{m1} (ary '{n1} (ary '{m} (ary '{n} num)))))  
    (dl2:exp(ary '{p1} (ary '{n1} (ary '{m} (ary '{n} num)))))
  := mapPar (block_l_mul_dl zero mzero dl2) dl1.

Definition block_add_vec {m n n1:nat}
    (dl1 dl2:exp( ary '{n1} (ary '{m} (ary '{n} num))))
    :exp (ary '{n1} (ary '{m} (ary '{n} num))) :=
  mapSeq (cur mat_add_ss) (zip dl1 dl2).

(* bA + bB *)
Definition block_add {m n m1 n1:nat}
    (dl1 dl2:exp(ary '{m1} (ary '{n1} (ary '{m} (ary '{n} num)))))
    :exp(ary '{m1} (ary '{n1} (ary '{m} (ary '{n} num)))) :=
  mapPar (cur block_add_vec) (zip dl1 dl2).

Hint Unfold dl_mul_dl_pss l_mul_dl_ss product_seq cur.

Section matrix_mul.
Variable m n p :nat.
Hypotheses wtm:writeNat m = "m".
Hypotheses wtn:writeNat n = "n".
Hypotheses wtp:writeNat p = "p".

Let ma : exp (ary '{m} (ary '{n} num)) :=  (vrd "ma").
Let mb : exp (ary '{p} (ary '{n} num)) :=  (vrd "mb").
Let zero : exp num := vrd "0".

(* A × B *)
Lemma example_dl_mul_dl : exists output,
  CodeGen_comm 
    (A{dl_mul_dl_pss zero ma mb, acc' (Var "out")} ) 0  = output.
Proof.
  unfold dl_mul_dl_pss,l_mul_dl_ss,product_seq,cur.
  Translate.
  exists "#pragma omp parallel for
for (int i0=0; i0 < m; i0 += 1) 
  { for(int i1=0; i1<p; i1 += 1) 
  { float v2[n];
for(int i3=0; i3<n; i3 += 1) 
  { v2[i3] = (ma[i3][i0] * mb[i3][i1]) ;}
float v4;
v4 = 0 ;
for(int i5=0; i5<n; i5 += 1) 
  { v4 = (v2[i5] + v4) ;}
out[i0][i1] = v4 ;}}". reflexivity.
  Qed.

End matrix_mul.

Section block_matrix_mul.

Variable m1 n1 p1 m n p:nat.
Hypotheses wtm:writeNat m = "m2".
Hypotheses wtn:writeNat n = "n2".
Hypotheses wtp:writeNat p = "p2".
Hypotheses wtm1:writeNat m1 = "m1".
Hypotheses wtn1:writeNat n1 = "n1".
Hypotheses wtp1:writeNat p1 = "p1".
Hypotheses mzero_ele: forall p,applyps "mzero" p = "0".

Let xs : exp (ary '{m1} (ary '{n1} (ary '{m} (ary '{n} num)))) := (vrd "xs").
Let ys : exp (ary '{p1} (ary '{n1} (ary '{m} (ary '{n} num)))) := (vrd "ys").
Let zero : exp num := vrd "zero".
Let mzero :exp (ary '{m} (ary '{m} num)) := vrd "mzero".

Lemma block_dl_mul_dl_codes :exists output,
  CodeGen_comm
    (A{ block_dl_mul_dl zero mzero xs ys , acc' (Var "out")}) 0 = output.
Proof.
  unfold block_dl_mul_dl,block_l_mul_dl,block_product,mat_add_ps,vec_add_seq.
  unfold dl_mul_dl_sss,l_mul_dl_ss,product_seq,cur.
  Translate.
  exists "#pragma omp parallel for
for (int i0=0; i0 < m1; i0 += 1) 
  { #pragma omp parallel for
for (int i1=0; i1 < p1; i1 += 1) 
  { float v2[n1][m2][m2];
for(int i3=0; i3<n1; i3 += 1) 
  { for(int i4=0; i4<m2; i4 += 1) 
  { for(int i5=0; i5<m2; i5 += 1) 
  { float v6[n2];
for(int i7=0; i7<n2; i7 += 1) 
  { v6[i7] = (xs[i7][i4][i3][i0] * ys[i7][i5][i3][i1]) ;}
float v8;
v8 = zero ;
for(int i9=0; i9<n2; i9 += 1) 
  { v8 = (v6[i9] + v8) ;}
v2[i3][i4][i5] = v8 ;}}}
float v10[m2][m2];
for(int i11=0; i11<m2; i11 += 1) 
  { for(int i12=0; i12<m2; i12 += 1) 
  { v10[i11][i12] = mzero[i12][i11] ;}}
for(int i11=0; i11<n1; i11 += 1) 
  { #pragma omp parallel for
for (int i12=0; i12 < m2; i12 += 1) 
  { for(int i13=0; i13<m2; i13 += 1) 
  { v10[i12][i13] = (v2[i13][i12][i11] + v10[i13][i12]) ;}}}
#pragma omp parallel for
for (int i14=0; i14 < m2; i14 += 1) 
  { for(int i15=0; i15<m2; i15 += 1) 
  { out[i0][i1][i14][i15] = v10[i15][i14] ;}}}}".
  reflexivity.
Qed.

End block_matrix_mul.


  













